﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using System.Windows.Forms;
using Ast.Core.Events;
using Ast.Core.Services;
using Ast.Core.Services.Extensions;
using Ast.Core.Settings;
using Ast.Core.WinApi.DbgHelp;
using Ast.Framework.Forms;
using Castle.Core.Logging;

namespace Ast
{
    [FormMetadata(FormType.Symbol, IsTabForm = true)]
    public partial class SymbolForm : Form, ISymbolEvents
    {
        private readonly IPsLoadedModuleListService _psLoadedModuleListService;
        private readonly ISymbolService _symbolService;
        private readonly IPeService _peService;
        private readonly IStructureConverter _structureConverter;
        private readonly IAstSettings _astSettings;

        private CancellationTokenSource _queryTokenSource;
        private CancellationTokenSource _maskTokenSource;

        public SymbolForm(IPsLoadedModuleListService psLoadedModuleListService, ISymbolService symbolService, IPeService peService, IStructureConverter structureConverter, IAstSettings astSettings)
        {
            _psLoadedModuleListService = psLoadedModuleListService;
            _symbolService = symbolService;
            _peService = peService;
            _structureConverter = structureConverter;
            _astSettings = astSettings;

            InitializeComponent();

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        private async void SymbolForm_Load(object sender, EventArgs e)
        {
            await Task.Run(() =>
            {
                Invoke(new Action(() =>
                {
                    // update symbol tag list
                    coboSymTag.BeginUpdate();
                    coboSymTag.Items.AddRange(Enum.GetValues(typeof(DbgHelpConstant.SymTagEnum)).Cast<object>().ToArray());
                    coboSymTag.SelectedItem = DbgHelpConstant.SymTagEnum.SymTagUdt;
                    coboSymTag.EndUpdate();

                    // update module list
                    lbModules.BeginUpdate();
                    lbModules.Items.Clear();
                    _symbolService.ListModules((moduleName, moduleBase) =>
                    {
                        var moduleViewItem = new ModuleViewItem {Name = moduleName, Base = moduleBase};
                        lbModules.Items.Add(moduleViewItem);
                    });
                    lbModules.EndUpdate();
                }));

                // load kernel module
                var ntModule = _psLoadedModuleListService.GetNtModule();
                var ntModuleName = ntModule != null ? ntModule.Name : "ntoskrnl.exe";
                var ntModuleBase = ntModule != null ? ntModule.StartAddress : 0;
                var baseOfDll = _symbolService.LoadModule(ntModuleName, null, ntModuleBase);
                Logger.InfoFormat("loaded kernel symbol at 0x{0:x}", baseOfDll);

                // perform query
                Invoke(new Action(() => PerformQuery("*!*")));
            }).ConfigureAwait(false);
        }

        private async void btnLoad_Click(object sender, EventArgs e)
        {
            if (openFileDialog1.ShowDialog() != DialogResult.OK) return;

            await Task.Run(() =>
            {
                foreach (var file in openFileDialog1.FileNames.Where(File.Exists))
                {
                    LoadModule(file);
                }
            }).ConfigureAwait(false);
        }

        private void LoadModule(string file)
        {
            // determine file size
            uint size;
            if (".pdb".Equals(Path.GetExtension(file), StringComparison.InvariantCultureIgnoreCase))
            {
                size = (uint) new FileInfo(file).Length;
            }
            else
            {
                var pe = _peService.Load(file);
                if (pe == null)
                {
                    Logger.ErrorFormat("cannot load as pe file {0}", file);
                    return;
                }
                size = pe.OptionalHeader.SizeOfImage;
            }

            // find address space
            var baseOfDll = _symbolService.FindAddressSpace(size);
            if (baseOfDll == 0)
            {
                Logger.Error("cannot find available address space for loading symbol");
                return;
            }

            // load module in found address space
            _symbolService.LoadModule(file, null, baseOfDll);
        }

        private async void btnQuery_Click(object sender, EventArgs e)
        {
            var mask = txbMask.Text;
            if (string.IsNullOrEmpty(mask)) return;

            // symTag
            var symTag = (DbgHelpConstant.SymTagEnum) coboSymTag.SelectedItem;

            // cancel token
            if (_queryTokenSource != null && !_queryTokenSource.IsCancellationRequested) _queryTokenSource.Cancel();
            var tokenSource = new CancellationTokenSource();
            _queryTokenSource = tokenSource;

            try
            {
                await Query(symTag, mask, tokenSource).ConfigureAwait(false);
            }
            catch (OperationCanceledException)
            {
                // canceled
            }
        }

        private async Task Query(DbgHelpConstant.SymTagEnum symTag, string mask, CancellationTokenSource tokenSource)
        {
            await Task.Run(() =>
            {
                var items = new List<object>();
                _symbolService.SearchSymbols(0, symTag, mask, info =>
                {
                    if (tokenSource.IsCancellationRequested) return false; // canceled
                    var symbolViewItem = new SymbolViewItem
                    {
                        ModuleBase = info.ModBase,
                        ModuleName = FindModuleName(info.ModBase),
                        Name = info.Name,
                    };
                    items.Add(symbolViewItem);
                    return true;
                });

                // canceled
                if (tokenSource.IsCancellationRequested)
                {
                    Logger.WarnFormat("canceled current query {0}", mask);
                    return;
                }

                // symbol limits
                Logger.InfoFormat("found {0} symbols", items.Count);
                if (items.Count > _astSettings.SymbolFormDisplayLimits)
                {
                    Logger.WarnFormat("only the first {0} symbols displayed because of too many symbols found", _astSettings.SymbolFormDisplayLimits);
                }

                // update view
                Invoke(new Action(() =>
                {
                    lbSymbols.BeginUpdate();
                    lbSymbols.Items.Clear();
                    lbSymbols.Items.AddRange(items.Take(_astSettings.SymbolFormDisplayLimits).ToArray());
                    lbSymbols.EndUpdate();
                }));
            }, tokenSource.Token).ConfigureAwait(false);
        }

        private void lbModules_SelectedIndexChanged(object sender, EventArgs e)
        {
            var moduleViewItem = lbModules.SelectedItem as ModuleViewItem;
            if (moduleViewItem == null) return;

            PerformQuery(string.Format("{0}!*", moduleViewItem.Name));
        }

        private async void btnClear_Click(object sender, EventArgs e)
        {
            await Task.Run(() => { _symbolService.UnloadAllModules(); }).ConfigureAwait(false);
            Invoke(new Action(() =>
            {
                lbSymbols.BeginUpdate();
                lbSymbols.Items.Clear();
                lbSymbols.EndUpdate();
            }));
        }

        private async void btnConvert_Click(object sender, EventArgs e)
        {
            var symbolViewItem = lbSymbols.SelectedItem as SymbolViewItem;
            if (symbolViewItem == null) return;

            await Task.Run(() =>
            {
                var structureStr = _structureConverter.Convert(symbolViewItem.ModuleBase, symbolViewItem.Name);
                if (string.IsNullOrEmpty(structureStr)) return;

                Logger.InfoFormat("\r\n{0}", structureStr);
                Invoke(new Action(() => Clipboard.SetText(structureStr)));
            }).ConfigureAwait(false);
        }

        private async void txbMask_TextChanged(object sender, EventArgs e)
        {
            // cancel token
            if (_queryTokenSource != null && !_maskTokenSource.IsCancellationRequested) _maskTokenSource.Cancel();
            var tokenSource = new CancellationTokenSource();
            _maskTokenSource = tokenSource;

            try
            {
                await Task.Delay(500, tokenSource.Token)
                    .ContinueWith(task => { Invoke(new Action(() => btnQuery.PerformClick())); }, TaskContinuationOptions.OnlyOnRanToCompletion)
                    .ConfigureAwait(false);
            }
            catch (OperationCanceledException)
            {
                // canceled
            }
        }

        private void lbSymbols_DoubleClick(object sender, EventArgs e)
        {
            btnConvert.PerformClick();
        }

        private void PerformQuery(string query)
        {
            txbMask.Text = query;
            txbMask.SelectionStart = txbMask.TextLength;
            txbMask.Focus();
        }

        private ModuleViewItem FindModule(ulong baseOfDll)
        {
            return lbModules.Items.Cast<ModuleViewItem>().FirstOrDefault(moduleViewItem => moduleViewItem.Base == baseOfDll);
        }

        private string FindModuleName(ulong baseOfDll)
        {
            var module = FindModule(baseOfDll);
            return module != null ? module.Name : null;
        }

        #region Events

        void ISymbolEvents.OnLoadedModule(string imageName, string moduleName, ulong baseOfDll)
        {
            Invoke(new Action(() =>
            {
                var moduleViewItem = new ModuleViewItem {Name = moduleName, Base = baseOfDll};
                lbModules.Items.Add(moduleViewItem);
            }));
        }

        void ISymbolEvents.OnUnloadedModule(ulong baseOfDll)
        {
            Invoke(new Action(() =>
            {
                var foundModule = FindModule(baseOfDll);
                if (foundModule != null)
                {
                    lbModules.Items.Remove(foundModule);
                }
            }));
        }

        #endregion

        public class ModuleViewItem
        {
            public string Name { get; set; }
            public ulong Base { get; set; }

            public override string ToString()
            {
                return string.Format("{0} 0x{1:x}", Name, Base);
            }
        }

        public class SymbolViewItem
        {
            public string ModuleName { get; set; }
            public ulong ModuleBase { get; set; }
            public string Name { get; set; }

            public override string ToString()
            {
                return string.Format("{0}!{1}", ModuleName, Name);
            }
        }
    }
}